package com.omni.monitor.plugin.threadpool;

import com.omni.monitor.model.threadpool.ThreadPoolStats;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.Assert;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;

/**
 * 支持自定义规则的线程池自动扩展器
 */
@Slf4j
public class ThreadPoolAutoScale {

    // 线程池监控器（用于获取线程池状态）
    private final ThreadPoolMonitor threadPoolMonitor;

    // 规则映射：线程池名称 -> 自定义调整规则
    private final Map<String, Function<ThreadPoolStats, Integer>> rules = new HashMap<>();

    // 检查间隔（默认5秒）
    private long checkInterval = 5000;

    // 定时任务执行器
    private ScheduledExecutorService scheduler;

    // 是否运行中
    private final AtomicBoolean running = new AtomicBoolean(false);

    public ThreadPoolAutoScale(ThreadPoolMonitor threadPoolMonitor) {
        Assert.notNull(threadPoolMonitor, "线程池监控器不能为空");
        this.threadPoolMonitor = threadPoolMonitor;
    }

    /**
     * 添加线程池调整规则
     * @param poolName 线程池名称
     * @param rule 规则：接收线程池统计信息，返回目标核心线程数
     */
    public ThreadPoolAutoScale addRule(String poolName, Function<ThreadPoolStats, Integer> rule) {
        Assert.hasText(poolName, "线程池名称不能为空");
        Assert.notNull(rule, "规则不能为空");
        rules.put(poolName, rule);
        return this;
    }

    /**
     * 启动自动调整
     */
    public void start() {
        if (running.compareAndSet(false, true)) {
            scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
                Thread t = new Thread(r, "thread-pool-auto-scaler");
                t.setDaemon(true);
                return t;
            });
            scheduler.scheduleAtFixedRate(this::adjustAllPools, 0, checkInterval, TimeUnit.MILLISECONDS);
            log.info("线程池自动扩展器启动，监控线程池数量: {}", rules.size());
        }
    }

    /**
     * 停止自动调整
     */
    public void stop() {
        if (running.compareAndSet(true, false) && scheduler != null) {
            scheduler.shutdown();
            log.info("线程池自动扩展器停止");
        }
    }

    /**
     * 核心逻辑：调整所有注册了规则的线程池
     */
    private void adjustAllPools() {
        try {
            // 遍历所有规则，应用调整
            for (Map.Entry<String, Function<ThreadPoolStats, Integer>> entry : rules.entrySet()) {
                String poolName = entry.getKey();
                Function<ThreadPoolStats, Integer> rule = entry.getValue();

                // 1. 获取线程池当前状态
                ThreadPoolStats stats = threadPoolMonitor.getThreadPoolStats(poolName);
                if (stats == null) {
                    log.warn("线程池[{}]未找到，跳过调整", poolName);
                    continue;
                }

                // 2. 应用规则计算目标核心线程数
                int targetCoreSize = rule.apply(stats);

                // 3. 边界校验（不超过线程池最大线程数，不小于1）
                targetCoreSize = Math.max(1, Math.min(targetCoreSize, stats.getMaximumPoolSize()));

                // 4. 执行调整（如果目标值与当前值不同）
                if (targetCoreSize != stats.getCorePoolSize()) {
                    threadPoolMonitor.setCorePoolSize(poolName, targetCoreSize);
                    log.info("线程池[{}]调整 - 核心线程数从 {} -> {}",
                            poolName, stats.getCorePoolSize(), targetCoreSize);
                }
            }
        } catch (Exception e) {
            log.error("线程池自动调整失败", e);
        }
    }

    // Setter方法（用于自定义检查间隔）
    public void setCheckInterval(long checkInterval) {
        this.checkInterval = checkInterval;
    }
}
